from typing import List
import streamlit as st
from langchain_core.documents.base import Document
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.runnables.base import Runnable
from langchain_core.runnables.utils import Output
from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate
from langchain.vectorstores import VectorStore


class ChatAgent:
    def __init__(self, prompt: ChatPromptTemplate, llm: Runnable):
        """
        初始化 ChatAgent。

        参数：
        - prompt (ChatPromptTemplate): 聊天提示模板。
        - llm (Runnable): 语言模型可运行实例。
        """
        self.history = StreamlitChatMessageHistory(key="chat_history")
        self.llm = llm
        self.prompt = prompt
        self.chain = self.setup_chain()
 
    def reset_history(self) -> None:
        """
        清除聊天历史记录，开始新的聊天会话。
        """
        self.history.clear()

    def setup_chain(self) -> RunnableWithMessageHistory:
        """
        设置 ChatAgent 的链条。

        返回：
        - RunnableWithMessageHistory: 配置好的链条，包含消息历史记录。
        """
        chain = self.prompt | self.llm
        return RunnableWithMessageHistory(
            chain,
            lambda session_id: self.history,
            input_messages_key="question",
            history_messages_key="history",
        )

    def display_messages(self, selected_query: str) -> None:
        """
        在聊天界面展示消息。
        如果没有历史消息，添加默认的 AI 消息。
        """
        if len(self.history.messages) == 0:
            self.history.add_ai_message(f"Let's chat about your query: {selected_query}")
        for msg in self.history.messages:
            st.chat_message(msg.type).write(msg.content)

    def format_retreieved_abstracts_for_prompt(self, documents: List[Document]) -> str:
        """
        格式化检索到的文档为字符串，传递给 LLM。
        """
        formatted_strings = []
        for doc in documents:
            formatted_str = f"ABSTRACT TITLE: {doc.metadata['title']}, ABSTRACT CONTENT: {doc.page_content}, ABSTRACT DOI: {doc.metadata['source'] if 'source' in doc.metadata.keys() else 'DOI missing..'}"
            formatted_strings.append(formatted_str)
        return "; ".join(formatted_strings)

    def get_answer_from_llm(self, question: str, retrieved_documents: List[Document]) -> Output:
        """
        根据用户问题和检索到的文档，从 LLM 获取响应。
        """
        config = {"configurable": {"session_id": "any"}}
        return self.chain.invoke(
            {
                "question": question, 
                "retrieved_abstracts": retrieved_documents,
            }, config
        )

    def retrieve_documents(self, retriever: VectorStore, question: str, cut_off: int = 5) -> List[Document]:
        """
        使用相似度搜索检索文档
        cut_off 参数控制检索结果的数量（默认值为 5）。
        """
        return retriever.similarity_search(question)[:cut_off]

    def start_conversation(self, retriever: VectorStore, selected_query: str) -> None:
        """
        在聊天界面开始对话。
        显示消息，提示用户输入，并处理 AI 的响应。
        """
        self.display_messages(selected_query)
        user_question = st.chat_input(placeholder="Ask me anything..")
        if user_question:
            documents = self.retrieve_documents(retriever, user_question)
            retrieved_abstracts = self.format_retreieved_abstracts_for_prompt(documents)
            st.chat_message("human").write(user_question)
            response = self.get_answer_from_llm(user_question, retrieved_abstracts)
            st.chat_message("ai").write(response.content)
